#include "Containers/RangePair.h"
#include "File/BOSSOutputLoader.h"
#include "File/YamlLoader.h"
#include "Fit/RooFitBuilder.h"
#include "Functions/Draw.h"
#include "Functions/Error.h"
#include "Functions/File.h"
#include "Functions/PrintTable.h"
#include "Functions/TerminalIO.h"
#include "Globals/Settings.h"
#include "Particle/ReconstructedParticle.h"
#include "RooFitResult.h"
#include "RooHistPdf.h"
#include "RooPlot.h"
#include "TArrow.h"
#include "TCanvas.h"
#include "TLatex.h"
#include "TLine.h"
#include <cmath>
#include <iostream>

using namespace BOSS_Afterburner;
using namespace BOSS_Afterburner::Error;
using namespace BOSS_Afterburner::File;
using namespace BOSS_Afterburner::TerminalIO;
using namespace RooFit;
using namespace BOSS_Afterburner::Detector::Units;
using namespace std;

void FitSigPlusBck(const char* config_file)
{
  YamlLoader       conf(config_file);
  BOSSOutputLoader file(conf);

  // * Cut ranges * //
  const auto& cutvals = conf["cutvals"];
  RangePair   cutD0   = cutvals["D0"];
  RangePair   cutphi  = cutvals["phi"];

  // * Plot and fit ranges * //
  const auto& ranges = conf["Ranges"];
  RangePair   fit    = ranges["fit"];
  RangePair   plot   = ranges["plot"];

  // * Construct fit variable * //
  ReconstructedParticle particle(conf["Reconstructed particle"]);
  const char*           varname =
    Form("#it{M}_{%s#rightarrow %s}", particle.NameLaTeX(), particle.DaughterLabel());
  RooRealVar    var(varname, varname, plot.X1(), plot.X2(), GeVc2);
  auto          frame = var.frame();
  RooFitBuilder builder(var);

  // * Construct RooFit data histogram * //
  auto        hist = file.Draw(conf["Data to fit"]);
  RooDataHist dataHist("data", hist->GetTitle(), var, Import(*hist));

  // * Determine signal model * //
  const auto& sigDef = conf["Signal model"];
  if(!sigDef || !sigDef.IsMap()) EXCEPTION("Missing key \"Signal model\"");
  const auto& sigTypeDef = sigDef["type"];
  if(!sigTypeDef || !sigTypeDef.IsScalar()) EXCEPTION("Missing key \"type\" in \"Signal model\"");
  RooAbsPdf* sigModel = nullptr;
  TString    sigType  = sigTypeDef.as<string>().c_str();
  if(sigType.EqualTo("Signal Monte Carlo"))
    sigModel = builder.BuildHistPdf(sigDef);
  else if(sigType.EqualTo("Voigtian"))
    sigModel = builder.BuildVoigtian(sigDef);
  else if(sigType.EqualTo("Breit-Wigner"))
    sigModel = builder.BuildBreitWigner(sigDef);
  else if(sigType.EqualTo("Gaussian"))
    sigModel = builder.BuildGaussian(sigDef);
  else if(sigType.EqualTo("P-wave"))
    sigModel = builder.BuildPwave(sigDef);
  else
    EXCEPTION("No valid signal model type");

  // * Determine background model * //
  const auto& bckDef = conf["Background model"];
  if(!bckDef || !bckDef.IsMap()) EXCEPTION("Missing key \"Background model\"");
  const auto& bckTypeDef = bckDef["type"];
  if(!bckTypeDef || !bckTypeDef.IsScalar())
    EXCEPTION("Missing key \"type\" in \"Background model\"");
  RooAbsPdf* bckModel = nullptr;
  TString    bckType  = bckTypeDef.as<string>().c_str();
  if(bckType.EqualTo("Argus"))
    bckModel = builder.BuildArgus(bckDef);
  else if(bckType.EqualTo("Argus BESIII"))
    bckModel = builder.BuildArgusBESIII(bckDef);
  else if(bckType.EqualTo("Truncated polynomial"))
    bckModel = builder.BuildTruncatedPolynomial(bckDef);
  else
    EXCEPTION("No valid background model type");

  // * Merge signal and background model * //
  const auto& normDef = conf["Normalisation factors"];
  if(!normDef || !normDef.IsMap()) EXCEPTION("Missing key \"Normalisation factors\"");
  auto       Nsig = builder.BuildVar(normDef["signal"]);
  auto       Nbck = builder.BuildVar(normDef["background"]);
  RooArgList sum(*sigModel, *bckModel);
  RooArgList coeff(*Nsig, *Nbck);
  auto       fullShape = builder.AddPdfs(&sum, &coeff);

  // * Perform fit * //
  Color_t colour = kRed;
  cout << "here0" << endl;
  fullShape->plotOn(frame, LineColor(colour), LineStyle(kDotted), Range(plot.X1(), plot.X2()),
                    LineWidth(1));
  cout << "here1" << endl;
  cout << fullShape << endl;
  auto result = fullShape->fitTo(dataHist, Range(fit.X1(), fit.X2()), Save(true));
  cout << "here2" << endl;
  fullShape->plotOn(frame, LineColor(colour), LineStyle(kDotted), Range(plot.X1(), plot.X2()),
                    LineWidth(1));
  cout << "here3" << endl;
  fullShape->plotOn(frame, LineColor(colour), LineWidth(1));
  cout << "here4" << endl;
  fullShape->plotOn(frame, Components(*bckModel), LineColor(colour), LineStyle(kDashed),
                    LineWidth(1));
  cout << "here5" << endl;
  fullShape->paramOn(frame, Layout(.68, .97, .92), Format("NU"));
  cout << "here6" << endl;

  // * Configure frame style * //
  if(frame)
  {
    frame->getAttText()->SetTextSize(.035);
    frame->getAttText()->SetTextColor(colour);
  }

  // * Draw and safe results * //
  TCanvas c;
  c.SetBatch();
  frame->Draw();
  Draw::IndicateFitRange(conf);
  Draw::Indicate3SigmaRange(conf);
  sigType.ReplaceAll(" ", "");
  bckType.ReplaceAll(" ", "");
  Draw::SaveCanvas(Form("%s_%s+%s", particle.Name(), sigType.Data(), bckType.Data()), &c);

  // * Fit results * //
  TerminalIO::PrintHeader("Fit results for first fit");
  cerr << "  Signal:           N_{sig} = " << Nsig->getVal() << endl;
  cerr << "  Background:       N_{bck} = " << Nbck->getVal() << endl;
  cerr << "  Chi square:  \\chi^2_{red} = " << frame->chiSquare(2) << endl;
  cerr << "  Likelihood:       -log(L) = " << result->minNll() << endl;
  cerr << "                         L  = " << exp(-0.5 * result->minNll()) << endl;
}

/// Main function that is called when executing the executable compiled using e.g. `g++`.
int main(int argc, char* argv[])
{
  TerminalIO::CheckMainArguments(argc, argv, 1);
  FitSigPlusBck(argv[1]);
  return 0;
}